import numpy as np
import os

from util.utils import mkdirs
from util.download import download_zip

har_url = ["https://archive.ics.uci.edu/ml/machine-learning-databases/00240/UCI%20HAR%20Dataset.zip"]
har_meta = "UCI HAR Dataset"

har_file_tree = {
    "train": {
        "x": "X_train.txt",
        "y": "y_train.txt",
        "Inertial Signals": {
            "body_acc_x": "body_acc_x_train.txt",
            "body_acc_y": "body_acc_y_train.txt",
            "body_acc_z": "body_acc_z_train.txt",
            "body_gyro_x": "body_gyro_x_train.txt",
            "body_gyro_y": "body_gyro_y_train.txt",
            "body_gyro_z": "body_gyro_z_train.txt",
            "total_acc_x": "total_acc_x_train.txt",
            "total_acc_y": "total_acc_y_train.txt",
            "total_acc_z": "total_acc_z_train.txt",
        },
        "subject": "subject_train.txt",
    },
    "test": {
        "x": "X_test.txt",
        "y": "y_test.txt",
        "Inertial Signals": {
            "body_acc_x": "body_acc_x_test.txt",
            "body_acc_y": "body_acc_y_test.txt",
            "body_acc_z": "body_acc_z_test.txt",
            "body_gyro_x": "body_gyro_x_test.txt",
            "body_gyro_y": "body_gyro_y_test.txt",
            "body_gyro_z": "body_gyro_z_test.txt",
            "total_acc_x": "total_acc_x_test.txt",
            "total_acc_y": "total_acc_y_test.txt",
            "total_acc_z": "total_acc_z_test.txt",
        },
        "subject": "subject_test.txt",
    },
    "activity_labels": "activity_labels.txt"
}

class HARDATASET(object):
    def __init__(self, dir_="data"):
        if not os.path.exists(os.path.join(dir_, har_meta + ".zip")):
            print("har data not found!")
            self.download_har(dir_)
        self.file_root = os.path.join(dir_, har_meta)
        self.label_dict = None

    def y2label(self):
        labels = self.get_data("activity_labels", dtype=str)
        if self.label_dict is None:
            self.label_dict = {int(con[0]): con[1] for con in labels}
        return self.label_dict

    def get_data(self, *path, dtype=np.float32):
        file_ = os.path.join(self.file_root, *path[0:-1])
        root = har_file_tree
        for p in path:
            if p in root:
                root = root[p]
            else:
                raise FileNotFoundError
        file_ = os.path.join(file_, root)
        return np.loadtxt(file_, dtype=dtype)

    def download_har(self, dir_):
        download_zip([os.path.join(dir_, har_meta + ".zip")], har_url)
